#include <iostream>
#include <cmath>
#include <algorithm>
#include <string>
#include <time.h>

#include <TROOT.h>
#include <TFile.h>
#include <TTree.h>
#include <TRandom.h>
#include <TMatrixD.h>
#include <TLorentzVector.h>

#include "GeomSvc.h"
#include "SRawEvent.h"
#include "KalmanUtil.h"
#include "KalmanTrack.h"
#include "KalmanFilter.h"
#include "KalmanFinder.h"
#include "KalmanFitter.h"
#include "VertexFit.h"
#include "SRecEvent.h"


using namespace std;

int main(int argc, char *argv[])
{
  //Initialize geometry service
  Log("Initializing geometry service ... ");
  GeomSvc* geometrySvc = GeomSvc::instance();
  geometrySvc->init(GEOMETRY_VERSION);

  //Retrieve the raw event
  Log("Retrieving the event stored in ROOT file ... ");

  SRawEvent *rawEvent = new SRawEvent();
  int nSeeds, nSeedsX, nSeedsY;
  double ax[50], ay[50], bx[50], by[50];
  int xIndex[50], yIndex[50];

  TFile *dataFile = new TFile(argv[1], "READ");
  TTree *dataTree = (TTree *)dataFile->Get("save");

  dataTree->SetBranchAddress("rawEvent", &rawEvent);
  dataTree->SetBranchAddress("nSeeds", &nSeeds);
  dataTree->SetBranchAddress("nSeedsX", &nSeedsX);
  dataTree->SetBranchAddress("nSeedsY", &nSeedsY);
  dataTree->SetBranchAddress("ax", ax);
  dataTree->SetBranchAddress("bx", bx);
  dataTree->SetBranchAddress("ay", ay);
  dataTree->SetBranchAddress("by", by);
  dataTree->SetBranchAddress("xIndex", xIndex);
  dataTree->SetBranchAddress("yIndex", yIndex);

  double time;
  int nCandidates;
  double chisq;
  int charge;
  int nRecHits;
  int iRecHits[50];
  int nHodoHits[3];
  int detectorID[50];
  double R[50], T[50], res[50];
  double x[50], y[50], z[50];
  double px[50], py[50], pz[50], p[50];

  double z_vertex, chisq_vertex;
  double p_vertex, px_vertex, py_vertex, pz_vertex;

  SRecEvent *recEvent = new SRecEvent();

  TFile *saveFile = new TFile(argv[2], "recreate");
  TTree *saveTree = dataTree->CloneTree(0);

  saveTree->Branch("recEvent", &recEvent, 256000, 99);

  saveTree->Branch("nCandidates", &nCandidates, "nCandidates/I");
  saveTree->Branch("time", &time, "time/D");
  saveTree->Branch("chisq", &chisq, "chisq/D");
  saveTree->Branch("charge", &charge, "charge/I");
  saveTree->Branch("nRecHits", &nRecHits, "nRecHits/I");
  saveTree->Branch("nHodoHits", nHodoHits, "nHodoHits[3]/I");
  saveTree->Branch("iRecHits", iRecHits, "iRecHits[nRecHits]/I");
  saveTree->Branch("dID", detectorID, "dID[nRecHits]/I");
  saveTree->Branch("res", res, "res[nRecHits]/D");
  saveTree->Branch("R", R, "R[nRecHits]/D");
  saveTree->Branch("T", T, "T[nRecHits]/D");
  saveTree->Branch("x", x, "x[nRecHits]/D");
  saveTree->Branch("y", y, "y[nRecHits]/D");
  saveTree->Branch("z", z, "z[nRecHits]/D");
  saveTree->Branch("px", px, "px[nRecHits]/D");
  saveTree->Branch("py", py, "py[nRecHits]/D");
  saveTree->Branch("pz", pz, "pz[nRecHits]/D");
  saveTree->Branch("p", p, "p[nRecHits]/D");

  saveTree->Branch("z_vertex", &z_vertex, "z_vertex/D");
  saveTree->Branch("chisq_vertex", &chisq_vertex, "chisq_vertex/D");
  saveTree->Branch("p_vertex", &p_vertex, "p_vertex/D");
  saveTree->Branch("px_vertex", &px_vertex, "px_vertex/D");
  saveTree->Branch("py_vertex", &py_vertex, "py_vertex/D");
  saveTree->Branch("pz_vertex", &pz_vertex, "pz_vertex/D");

  //Initialize track finder
  Log("Initializing the track finder and kalman filter ... ");
  KalmanFilter *filter = new KalmanFilter();
  KalmanFinder *finder = new KalmanFinder();
  VertexFit    *vtxfit = new VertexFit();

  //finder->initEvaluation("eval.root");

  int offset = argc > 3 ? atoi(argv[3]) : 0;
  int nEvtMax = argc > 4 ? atoi(argv[4]) + offset : dataTree->GetEntries();
  if(nEvtMax > dataTree->GetEntries()) nEvtMax = dataTree->GetEntries();
  Log("Running from event " << offset << " through to event " << nEvtMax);
  for(int i = offset; i < nEvtMax; i++)
    {
      dataTree->GetEntry(i);
      Log("Processing event " << i << " with eventID = " << rawEvent->getEventID());

      rawEvent->reIndex("oah");
      finder->setEvent(rawEvent);
      
      Seed _seed;
      _seed.axz = ax[xIndex[0]]; 
      _seed.ayz = ay[yIndex[0]];
      _seed.bxz = bx[xIndex[0]];
      _seed.byz = by[yIndex[0]];
      _seed.decideCharge();
      _seed.decideMomentum();

      ///Timing stack.getPositions(3, x, y, z);rts ...
      clock_t time_single = clock();

      nCandidates = finder->processOneSeed(_seed);
      KalmanTrack _track = finder->getBestCandidate();

      if(!_track.isValid()) continue;

      charge = _track.getCharge();
      nRecHits = _track.getPositions(3, x, y, z);
      _track.getAlignment(3, detectorID, res, R, T);
      _track.getMomentums(3, px, py, pz);
      _track.getHitsIndex(iRecHits);

      for(int j = 0; j < 3; j++)
	{
	  nHodoHits[j] = finder->getNHodoHits(j+1, _track);
	}

      for(int j = 0; j < nRecHits; j++)
	{
	  p[j] = sqrt(px[j]*px[j] + py[j]*py[j] + pz[j]*pz[j]);
	}

      chisq = _track.getChisq();

      time_single = clock() - time_single;
      time = double(time_single)/CLOCKS_PER_SEC;
      Log("It takes " << time << " seconds for this event.");
      _track.print();

      ///Fill the SRecEvent
      recEvent->setRawEvent(rawEvent);
      SRecTrack recTrack = _track.getSRecTrack();
      recTrack.setHodoHits(nHodoHits);
      recEvent->insertTrack(recTrack);

      z_vertex = vtxfit->findSingleMuonVertex(_track.getNodeList().front());
      p_vertex = _track.getMomentumVertex(z_vertex, px_vertex, py_vertex, pz_vertex);
      chisq_vertex = _track.getChisqVertex();

      ///Start the track vertex fit
      recEvent->reIndex();
      saveTree->Fill();

      recEvent->clear();
      rawEvent->clear();
    }

  saveFile->cd();
  saveTree->Write();
  saveFile->Close();

  filter->close();
  delete vtxfit;
  delete finder;

  return 1;
}
